# coding=utf-8
import os
import argparse
import xml.etree.ElementTree as ET
from TemplateXmlAndJsonIO import *

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='parsing XML to xml or json')
    parser.add_argument('--src', type=str, default='/home/Dukto/101_whole_passport_1.xml')
    parser.add_argument('--dst', type=str, default='/home/eric/Disk100G/githubProject/cvat/json_save_111')
    args = parser.parse_args()
    return args

def get_xml_data(save_path, filename = None):
    if filename is None:
        print("The file is empty, please ensure it exits!")
    else:
        tree = ET.parse(filename)
        root = tree.getroot()

        annotator = ''
        for ann in root.iter('owner'):
            annotator = ann.find('username').text

        for images in root.findall('image'):
            name = images.get('name')
            image_name = os.path.basename(name)
            path_name = os.path.dirname(name)
            width  = images.get('width')
            height = images.get('height')
            image_info = {'image_name': image_name, 'width': width, 'height': height, 'depth': 3, 'annotator': annotator}
            ### xml to xml
            if images.findall('box'):
                print("image %s has %d bbox." %(images.get('id'), len(images.findall('box'))))
                writerj = TemplateXmlAndJsonWriter(image_info)
                bbox_results = []
                for boxes in images.findall('box'):
                    cls = boxes.get('label')
                    xmin = int(float(boxes.get('xtl')))
                    ymin = int(float(boxes.get('ytl')))
                    xmax = int(float(boxes.get('xbr')))
                    ymax = int(float(boxes.get('ybr')))
                    bbox = [xmin, ymin, xmax, ymax]
                    box_info = {'class': cls, 'bbox': bbox}
                    bbox_results.append(box_info)
                xml_name = os.path.splitext(image_name)[0] + '.xml'
                dir_path = os.path.join(save_path, path_name)
                if not os.path.exists(dir_path):
                    os.mkdir(dir_path)
                else:
                    pass
                xml_save_path = os.path.join(dir_path, xml_name)
                writerj.save_xml(bbox_results, xml_save_path)

            ### xml to json
            elif images.findall('polygon'):
                print("image %s has %d polygon." %(images.get('id'), len(images.findall('polygon'))))
                writerj = TemplateXmlAndJsonWriter(image_info, shape_type="POLYGON")
                polygon_results = []
                id = 0
                for polygons in images.findall('polygon'):
                    points_list = []
                    cls = polygons.get('label')
                    points = polygons.get('points')
                    pp = points.split(";")
                    for i in range(len(pp)):
                        points_split = ((float(pp[i].split(",")[0]), float(pp[i].split(",")[1])))
                        points_list.append(points_split)
                        # print(points_list)
                    writerj.addPolygon(points_list, cls, annotator, id)
                    id += 1
                json_name = os.path.splitext(image_name)[0] + '.json'
                dir_path = os.path.join(save_path, path_name)
                if not os.path.exists(dir_path):
                    os.mkdir(dir_path)
                else:
                    pass
                json_save_path = os.path.join(dir_path, json_name)
                writerj.save_json(json_save_path)


if __name__ == '__main__':
    args = parse_args()
    print('Called with args:')
    print(args)
    get_xml_data(args.dst, args.src)